{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/NeuromatchAcademy/course-content/blob/W2D3%2C5/tutorials/W2D3_DecisionMaking/student/W2D3_Tutorial2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "# Neuromatch Academy: Week 2, Day 3, Tutorial 2\n",
    "\n",
    "# Hidden Markov Model\n",
    "\n",
    "__Content creators:__ Yicheng Fei with help from Jesse Livezey\n",
    "\n",
    "__Content reviewers:__ John Butler, Matt Krause, Meenakshi Khosla, Spiros Chavlis, Michael Waskom\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "# Tutorial objectives\n",
    "\n",
    "The world around us is often changing state over time, but we may only have access to these states through noisy sensory measurements. Similarly, organisms and neural systems often are thought to transition between a set of discrete states (up/down states, sleep/wake, etc.) which may only be indirectly observable through their impact on neural activity. Hidden Markov Models are a class of models that allow us to reason about the dynamics of a set of unobserved states that lead to the changing sensory inputs or data we observe.\n",
    "\n",
    "In this notebook, we'll first simulate a Hidden Markov Model and observe how changing the transition probability and observation noise impact what the samples look like. Then we'll look at how uncertainty increases as we make future predictions without evidence (from observations) and how to gain information from the observations.\n",
    "The HMM model we use in the first part of the tutorial will have binary latent variable $x_t \\in \\{0,1\\}$ and 1D Gaussian emmision model $y_t|x_t \\sim \\mathcal{N}(\\mu_{x_t},\\sigma^2_{x_t})$. You will learn how to:\n",
    "\n",
    "* Build an HMM in Python and generate sample data.\n",
    "* Calculate how predictive probabilities propagates in a Markov Chain with no evidence.\n",
    "* Combine new evidence and prediction from past evidence  to estimate latent states.\n",
    "\n",
    "---\n",
    "\n",
    "There is an an **optional** part for you to get a sense of how to perform parameter estimation of an HMM using the EM algorithm. **Please make sure to complete and understand the main exercises before proceeding.**\n",
    "\n",
    "In the optional part, you will implement an HMM of a network of Poisson spiking neurons mentioned in today's intro and:\n",
    "\n",
    "* Implement the forward-backward algorithm\n",
    "* Complete the E-step and M-step \n",
    "* Learning parameters for the example problem using the EM algorithm\n",
    "* Get an intuition of how the EM algorithm monotonically increase data likelihood"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 518
    },
    "colab_type": "code",
    "outputId": "0876f17e-ad9e-45d9-ca0b-2748204138f8"
   },
   "outputs": [],
   "source": [
    "#@title Video 1: Introduction\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"ceQXN0OUaFo\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "colab": {},
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "!pip install hmmlearn --quiet\n",
    "\n",
    "import numpy as np\n",
    "from scipy import stats\n",
    "from scipy.optimize import linear_sum_assignment\n",
    "from hmmlearn import hmm\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import patches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {},
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "#@title Figure Settings\n",
    "import ipywidgets as widgets       # interactive display\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "plt.style.use(\"https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {},
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "#@title Helper functions\n",
    "\n",
    "def plot_hmm1(model, states, observations):\n",
    "  \"\"\"Plots HMM states and observations for 1d states and observations.\n",
    "\n",
    "  Args:\n",
    "    model (hmmlearn model):               hmmlearn model used to get state means.\n",
    "    states (numpy array of floats):       Samples of the states.\n",
    "    observations (numpy array of floats): Samples of the states.\n",
    "  \"\"\"\n",
    "\n",
    "  nsteps = states.size\n",
    "  fig, ax1 = plt.subplots()\n",
    "  states_forplot = list(map(lambda s: model.means_[s], states))\n",
    "  ax1.step(np.arange(nstep), states_forplot, \"--\", where=\"mid\", alpha=1.0, c=\"green\")\n",
    "  ax1.set_xlabel(\"Time\")\n",
    "  ax1.set_ylabel(\"Latent State\", c=\"green\")\n",
    "  ax1.set_yticks([-1, 1])\n",
    "  ax1.set_yticklabels([\"State 1\", \"State 0\"])\n",
    "\n",
    "  ax2 = ax1.twinx()\n",
    "  ax2.plot(np.arange(nstep), observations.flatten(), c=\"blue\")\n",
    "  ax2.set_ylabel(\"Observations\", c=\"blue\")\n",
    "  ax1.set_ylim(ax2.get_ylim())\n",
    "  plt.show(fig)\n",
    "\n",
    "\n",
    "def plot_marginal_seq(predictive_probs, switch_prob):\n",
    "  \"\"\"Plots the sequence of marginal predictive distributions.\n",
    "\n",
    "    Args:\n",
    "      predictive_probs (list of numpy vectors): sequence of predictive probability vectors\n",
    "      switch_prob (float):                      Probability of switching states.\n",
    "  \"\"\"\n",
    "  T = len(predictive_probs)\n",
    "  prob_0 = [p_vec[0] for p_vec in predictive_probs]\n",
    "  prob_1 = [p_vec[1] for p_vec in predictive_probs]\n",
    "  fig, ax = plt.subplots()\n",
    "  ax.plot(np.arange(T), prob_0, color=\"orange\")\n",
    "  ax.plot(np.arange(T), prob_1, color=\"blue\")\n",
    "  ax.legend([\n",
    "    \"prob in state 0\", \"prob in state 1\"\n",
    "  ])\n",
    "  ax.text(T/2, 0.05, \"switching probability={}\".format(switch_prob), fontsize=12,\n",
    "          bbox=dict(boxstyle=\"round\", facecolor=\"wheat\", alpha=0.6))\n",
    "  ax.set_xlabel(\"Time\")\n",
    "  ax.set_ylabel(\"Probability\")\n",
    "  ax.set_title(\"Forgetting curve in a changing world\")\n",
    "  plt.show(fig)\n",
    "\n",
    "def simulate_prediction_only(model, nstep):\n",
    "  \"\"\"\n",
    "  Simulate the diffusion of HMM with no observations\n",
    "\n",
    "  Args:\n",
    "    model (hmm.GaussianHMM instance): the HMM instance\n",
    "    nstep (int): total number of time steps to simulate(include initial time)\n",
    "\n",
    "  Returns:\n",
    "    predictive_probs (list of numpy vector): the list of marginal probabilities\n",
    "  \"\"\"\n",
    "  entropy_list = []\n",
    "  predictive_probs = []\n",
    "  prob = model.startprob_\n",
    "  for i in range(nstep):\n",
    "    # calculate entropy\n",
    "    predictive_probs.append(prob)\n",
    "    # one step forward\n",
    "    prob = markov_forward(prob, model.transmat_)\n",
    "  return predictive_probs\n",
    "\n",
    "def plot_evidence_vs_noevidence(posterior_matrix, predictive_probs):\n",
    "  \"\"\"Plots the average posterior probabilities with evidence v.s. no evidence\n",
    "\n",
    "  Args:\n",
    "    posterior_matrix: (2d numpy array of floats): The posterior probabilities in state 1 from evidence (samples, time)\n",
    "    predictive_probs (numpy array of floats):  Predictive probabilities in state 1 without evidence\n",
    "  \"\"\"\n",
    "  nsample, T = posterior_matrix.shape\n",
    "  posterior_mean = posterior_matrix.mean(axis=0)\n",
    "  fig, ax = plt.subplots(1)\n",
    "  ax.plot([0.0, T],[0.5, 0.5], color=\"red\", linestyle=\"dashed\")\n",
    "  ax.plot(np.arange(T), predictive_probs, c=\"orange\", linewidth=2, label=\"No evidence\")\n",
    "  ax.scatter(np.tile(np.arange(T), (nsample, 1)), posterior_matrix, s=0.8, c=\"green\", alpha=0.3, label=\"With evidence(Sample)\")\n",
    "  ax.plot(np.arange(T), posterior_mean, c='green', linewidth=2, label=\"With evidence(Average)\")\n",
    "  ax.legend()\n",
    "  ax.set_yticks([0.0, 0.25, 0.5, 0.75, 1.0])\n",
    "  ax.set_xlabel(\"Time\")\n",
    "  ax.set_ylabel(\"Probability in State 0\")\n",
    "  ax.set_title(\"Gain confidence with evidence\")\n",
    "  plt.show(fig)\n",
    "\n",
    "def simulate_forward_inference(model, T, data=None):\n",
    "  \"\"\"\n",
    "  Given HMM `model`, calculate posterior marginal predictions of x_t for T-1 time steps ahead based on\n",
    "  evidence `data`. If `data` is not give, generate a sequence of observations from first component.\n",
    "\n",
    "  Args:\n",
    "    model (GaussianHMM instance): the HMM\n",
    "    T (int): length of returned array\n",
    "\n",
    "  Returns:\n",
    "    predictive_state1: predictive probabilities in first state w.r.t no evidence\n",
    "    posterior_state1: posterior probabilities in first state w.r.t evidence\n",
    "  \"\"\"\n",
    "\n",
    "  # First re-calculate hte predictive probabilities without evidence\n",
    "  predictive_probs = simulate_prediction_only(model, T)\n",
    "  # Generate an observation trajectory condtioned on that latent state x is always 1\n",
    "  if data is not None:\n",
    "    Y = data\n",
    "  else:\n",
    "    Y = np.asarray([model._generate_sample_from_state(0) for _ in range(T)])\n",
    "  # Calculate marginal for each latent state x_t\n",
    "  pt = np.exp(model._compute_log_likelihood(Y[[0]])) * model.startprob_\n",
    "  pt /= np.sum(pt)\n",
    "  posterior_probs = np.zeros((T, pt.size))\n",
    "  posterior_probs[0] = pt\n",
    "\n",
    "  for t in range(1, T):\n",
    "    posterior = one_step_update(model, posterior_probs[t-1], Y[[t]])\n",
    "    # normalize and add to the list\n",
    "    posterior /= np.sum(posterior)\n",
    "    posterior_probs[t] = posterior\n",
    "  posterior_state1 = np.asarray([p[0] for p in posterior_probs])\n",
    "  predictive_state1 = np.asarray([p[0] for p in predictive_probs])\n",
    "  return predictive_state1, posterior_state1\n",
    "\n",
    "def plot_forward_inference(model, states, observations, states_inferred):\n",
    "    \"\"\"Plot ground truth state sequence with noisy observations, and ground truth states v.s. inferred ones\n",
    "\n",
    "        Args:\n",
    "            model (instance of hmmlearn.GaussianHMM): an instance of HMM\n",
    "            states (numpy vector): vector of 0 or 1(int or Bool), the sequences of true latent states\n",
    "            observations (numpy vector of numpy vector): the un-flattened Gaussian observations at each time point, element has size (1,)\n",
    "            states_inferred (numpy vector): vector of 0 or 1(int or Bool), the sequences of inferred latent states\n",
    "    \"\"\"\n",
    "    plot_hmm1(model, states, observations)\n",
    "    fig, ax = plt.subplots()\n",
    "    # state 0 has larger mean\n",
    "    ax.step(np.arange(nstep), 1-states, color=\"green\", label=\"Ground Truth\")\n",
    "    ax.step(np.arange(nstep), 1-states_inferred, linestyle=\"dashed\", color=\"orange\", label=\"Inferred\")\n",
    "    ax.legend()\n",
    "    ax.set_title(\"Infer latent states from data\")\n",
    "    ax.set_xlabel(\"Time\")\n",
    "    ax.set_ylabel(\"Latent State\")\n",
    "    ax.set_yticks([0, 1])\n",
    "    ax.set_yticklabels([\"State 1\", \"State 0\"])\n",
    "    plt.show(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "# Section 1: Binary HMM with Gaussian observations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 518
    },
    "colab_type": "code",
    "outputId": "c4949ad8-5ef8-4cb7-e7f0-09cff2f9ce88"
   },
   "outputs": [],
   "source": [
    "#@title Video 2: Simulating a binary HMM with Gaussian observations\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"7cTnoe6Xt80\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### EXERCISE 1: Simulate a binary HMM with Gaussian observations\n",
    "\n",
    "In contrast with the Sequential Probability Ratio Test, the latent state in HMMs is no longer fixed over time. Instead, it can probabilistically switch or jump to a different state at each time step. However, the  time dependence of states at different times is simple: the probability of the state at time $t$ is totally determined by the state at time $t-1$. This is called called the **Markov property** or **Markovian** and the dependency of the whole state sequence $\\{x_1,...,x_t\\}$ can be described by a chain structure called a Markov Chain:\n",
    "\n",
    "<img src=\" https://github.com/NeuromatchAcademy/course-content/blob/master/tutorials/W2D3_DecisionMaking/static/W2D3_Tutorial2_markov_chain_diagram.png?raw=true\" alt=\"Markov chain drawing\" width=\"400\"/>\n",
    "\n",
    "Quantitatively, assuming we are at state $i$ at time $t-1$, and denote the probability to jump to state $j$ at next time point $t$ as $A_{ij}$, then the distribution at time $t$ can be calculated using only probabilities at $t-1$ and the transition matrix $A$:\n",
    "\n",
    "$$ p(x_t=j) = \\sum_j A_{ij} p(x_{t-1}=i) \\tag{1}$$\n",
    "\n",
    "$A_{ij}$ represents the transition probability to switch **FROM state $i$ TO state $j$** at next time step. Please note that this is contrast to the meaning used in the intro.\n",
    "\n",
    "We cannot directly observe the latent states $x_t$. What we can observe instead is a noisy measurement $y_t$ generated from $x_t$. Common choices of measurement/emission models include Gaussian distribution $y_t|x_t \\sim \\mathcal{N}(\\mu_{x_t},\\sigma^2_{x_t})$ for continuous observables, and Poisson distribution $y_t|x_t \\sim \\text{Pois}(\\lambda)$ for discrete observables (like spike counts).\n",
    "\n",
    "---\n",
    "\n",
    "In this exercise, you will use the package `hmmlearn` to implement a two-state HMM with Gaussian emissions. Your HMM will start in State 0 and transition between states (both $0 \\rightarrow 1$ and $1 \\rightarrow 0$) with probability `switch_prob`. Each state emits observations drawn from a Gaussian with $\\mu = 1$ for State 0 and $\\mu = -1$ for State 1. The variance of both states is fixed at `noise_level`.\n",
    "\n",
    "Please familarize yourself with the code and complete the following exercises in the next cell. You will need to:\n",
    "\n",
    "1. To implement the state transitions, complete the transition matrix  `transmat_` (i.e., $A_{ij}$) in the code below. \n",
    "\\begin{equation*}\n",
    "A_{i,j} = \n",
    "\\begin{pmatrix}\n",
    "p_{stay} & p_{switch} \\\\\n",
    "p_{switch} & p_{stay} \\\\\n",
    "\\end{pmatrix}\n",
    "\\end{equation*}\n",
    "with $p_{stay} = 1 - p_{switch}$. \n",
    "\n",
    "2. The *hidden* part of HMM means that we do not directly output the current state $x_t$, but instead observe a noisy emission $y_t | x_t$, here generated by a Gaussian. The  means have already been filled in for you, but you must complete the covariance matrix `covars_`. Set each state's observation variance to `noise_level`. The required shape, given below, is unusual, but don't overthink finding the (co)variance of a univariate Gaussian. \n",
    "\n",
    "3. Use the provided code to run and plot one realization of the HMM's output. \n",
    "4. Try different transition (`switch_prob`) and noise levels (`noise_level`) and observe their impact on the samples. You can do this by hand or use the widget in section **Interactive Cell**.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "def create_model(switch_prob=0.1, noise_level=1e-8, startprob=[1.0, 0.0]):\n",
    "    \"\"\"Create an HMM with binary state variable and 1D Gaussian observations\n",
    "    The probability to switch to the other state is `switch_prob`. Two\n",
    "    observation models have mean 1.0 and -1.0 respectively. `noise_level`\n",
    "    specifies the standard deviation of the observation models.\n",
    "\n",
    "    Args:\n",
    "        switch_prob (float): probability to jump to the other state\n",
    "        noise_level (float): standard deviation of observation models. Same for\n",
    "        two components\n",
    "\n",
    "    Returns:\n",
    "        model (hmm.GaussianHMM instance): the described HMM\n",
    "    \"\"\"\n",
    "    ############################################################################\n",
    "    # Insert your code here to:\n",
    "    #    * Create the transition matrix, `transmat_` so that the odds of\n",
    "    #      switching is `switch_prob`\n",
    "    #\t\t* Set the observation model variances, `covars_`, to `noise_level`\n",
    "    raise NotImplementedError(\"`create_model` is incomplete\")\n",
    "    ############################################################################\n",
    "    n_components = 2\n",
    "    model = hmm.GaussianHMM(n_components=n_components, covariance_type=\"full\")\n",
    "    model.startprob_ = np.asarray(startprob)\n",
    "    # should be shape (2, 2), i.e., a transition matrix for 2 states\n",
    "    model.transmat_ = ...\n",
    "    model.means_ = np.array([[1.0], [-1.0]])\n",
    "\n",
    "    # should be shape (2, 1, 1), i.e., 2 1x1 covariance matrices\n",
    "    model.covars_ = ...\n",
    "    model.sample(1)\n",
    "    return model\n",
    "\n",
    "\n",
    "np.random.seed(101)\n",
    "nstep = 50\n",
    "###############################################################################\n",
    "## After finishing the above exercises, please un-comment the following lines\n",
    "###############################################################################\n",
    "# model = create_model()\n",
    "# observations, states = model.sample(nstep)\n",
    "#plot_hmm1(model, states, observations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 433
    },
    "colab_type": "text",
    "outputId": "dbb97cf5-fc83-42c2-d418-0fd840f4c623"
   },
   "source": [
    "[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D3_DecisionMaking/solutions/W2D3_Tutorial2_Solution_75397371.py)\n",
    "\n",
    "*Example output:*\n",
    "\n",
    "<img alt='Solution hint' align='left' width=560 height=416 src=https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/tutorials/W2D3_DecisionMaking/static/W2D3_Tutorial2_Solution_75397371_0.png>\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 493,
     "referenced_widgets": [
      "ed6a8945ebe64f48ae2f3302f2e93206",
      "8ef2cc1cf8bb4c47afee79ab2a8cf992",
      "3ca24589e09a423a9b1fb3c901d8abbc",
      "c7095b8a311d4733bcc67b51488a4b71",
      "e2240f528edb4996b43495cb0b988add",
      "16adc5f0831a4350a179ecb8cc471a62",
      "d2c89481cf104bcfb44f70ec87383910",
      "7582992f3c1f4421bfcf39902ab353c4",
      "cd89f6bae56740969f6d4af4d43007be",
      "60271dc15a194e0d9cd6cbd37aa5d41e"
     ]
    },
    "colab_type": "code",
    "outputId": "7e3be85c-dc8b-4983-d02b-382f0070d3ec"
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "\n",
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "np.random.seed(101)\n",
    "nstep = 100\n",
    "\n",
    "@widgets.interact\n",
    "def plot(switch_prob=(0., 1, .01), log10_noise_level=(-8., 1., .01)):\n",
    "  model = create_model(switch_prob=switch_prob,\n",
    "                     noise_level=10.**log10_noise_level)\n",
    "\n",
    "  observations, states = model.sample(nstep)\n",
    "  observations = observations.flatten()\n",
    "  plot_hmm1(model, states, observations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "\n",
    "# Section 2: Forgetting information and gaining confidence with a known initial state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 518
    },
    "colab_type": "code",
    "outputId": "1ccda3e7-50af-473c-9f66-f9289e25528e"
   },
   "outputs": [],
   "source": [
    "#@title Video 3: Forgetting in a changing world\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"pRRo_L-n8nc\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### EXERCISE 2: Forgetting in a changing world\n",
    "\n",
    "Since the world (latent state) is changing over time, even if we know for sure that we are in state 0 at some time, we will be more and more uncertain that we'll remain in state 0 as time goes. In other words, when we try to make predictions of future states in a Markov Chain based on our current knowledge without future evidence, the influence of current state will decay over time.\n",
    "\n",
    "In this exercise, we'll inspect how we \"forget\" the current state information when predicting future states without any observation.\n",
    "\n",
    "Using the model you just defined, let's now make some predictions about $x_t$ given that we know $x_0=0$ for sure. We've already imposed this assumption by setting prior probabilities of $p(x_0)$ to $[1,0]$ earlier. \n",
    "\n",
    "**Suggestions** \n",
    "\n",
    "\n",
    "1. Complete the code in function `markov_forward` to calculate the predictive marginal distribution at next time step using `p_next = A.T @ p_current` \n",
    "\n",
    "3. Take a look at function `simulate_prediction_only` and understand how the predictive distribution propagates along the Markov chain\n",
    "\n",
    "4. Using our provided code, plot the predictive probabilities as a function of time\n",
    "\n",
    "5. Try different values of switching probability (`prob_switch`) either by hand or the widget in section **Interactive Cell** and observe how that affects the forgetting curve of current state. How does the curve look at when `prob_switch` $>0.5$?\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "def markov_forward(p0, A):\n",
    "  \"\"\"Calculate the forward predictive distribution in a discrete Markov chain\n",
    "\n",
    "  Args:\n",
    "    p0 (numpy vector): a discrete probability vector\n",
    "    A (numpy matrix): the transition matrix, A[i,j] means the prob. to\n",
    "    switch FROM i TO j\n",
    "\n",
    "  Returns:\n",
    "    p1 (numpy vector): the predictive probabilities in next time step\n",
    "  \"\"\"\n",
    "  ############################################################################\n",
    "  # Insert your code here to:\n",
    "  #      Compute the marginal distribution of Markov chain in next time step\n",
    "  #      Hint: use matrix multiply and be careful about the index orders\n",
    "  raise NotImplementedError(\"function `markov_forward` incomplete\")\n",
    "  ############################################################################\n",
    "  p1 = ...\n",
    "  return p1\n",
    "\n",
    "\n",
    "np.random.seed(101)\n",
    "T = 100\n",
    "switch_prob = 0.1\n",
    "noise_level = 2.0\n",
    "################################################################################\n",
    "# After finishing the above exercises, please un-comment the following lines\n",
    "################################################################################\n",
    "# model = create_model(switch_prob=switch_prob, noise_level=noise_level)\n",
    "# predictive_probs = simulate_prediction_only(model, T)\n",
    "\n",
    "#plot_marginal_seq(predictive_probs, switch_prob)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 433
    },
    "colab_type": "text",
    "outputId": "2f00190d-3387-4247-a146-0e537c564be0"
   },
   "source": [
    "[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D3_DecisionMaking/solutions/W2D3_Tutorial2_Solution_533408ce.py)\n",
    "\n",
    "*Example output:*\n",
    "\n",
    "<img alt='Solution hint' align='left' width=560 height=416 src=https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/tutorials/W2D3_DecisionMaking/static/W2D3_Tutorial2_Solution_533408ce_0.png>\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "- In the following interactive cell, play with different switching probabilities with the slider and observe how that affects the information loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 461,
     "referenced_widgets": [
      "0fecedf8473c4dea9b2a7defe16d04a2",
      "2d7ac9fb328242c48f8df962bc1cb366",
      "19066a381f42424e828d001d41621ca7",
      "2afec61b8cea4780bb4c99688d404455",
      "35952e0c5b30424dbfde6999be781115",
      "34fdf12432fe4e4391952fc24fde0bdc",
      "35c7300a00de432094c229426eb5d2b0"
     ]
    },
    "colab_type": "code",
    "outputId": "37218996-3f0a-4f62-8255-ac56d96d8e07"
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "\n",
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "np.random.seed(101)\n",
    "T = 100\n",
    "noise_level = 0.5\n",
    "\n",
    "@widgets.interact\n",
    "def plot(switch_prob=(0.01, .99, .01)):\n",
    "  model = create_model(switch_prob=switch_prob, noise_level=noise_level)\n",
    "  predictive_probs = simulate_prediction_only(model, T)\n",
    "  plot_marginal_seq(predictive_probs, switch_prob)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### EXERCISE 3: Gain confidence from evidence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 518
    },
    "colab_type": "code",
    "outputId": "36ead1a5-8750-4276-ca47-bd83749ccef6"
   },
   "outputs": [],
   "source": [
    "#@title Video 4: Gain confidence from evidence\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"dDjoxUxMgC0\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "As shown in Exercise 2, you lose information/gain uncertainty exponentially when predicting into future without further evidence because the Markov Chain in HMM will diffuse through a exponential number of possible paths. However HMM generates an observation at each time step and we can use these as evidence to improve our state estimate.\n",
    "\n",
    "<img src=\"https://github.com/NeuromatchAcademy/course-content/blob/master/tutorials/W2D3_DecisionMaking/static/W2D3_Tutorial2_HMM_diagram.png?raw=true\" alt=\"HMM drawing\" width=\"400\"/>\n",
    "\n",
    "Now let's incorporate evidence into our inference. In this exercise we will calculate the **posterior marginal** distribution $p(x_t|y_{1:t})$ recursively when evidence $y_t$ comes in at each time using the forward pass of the forward-backward algorithm.\n",
    "That is saying, suppose we know the posterior marginal of $x_{t-1}$ given all observations up to $t-1$, when a new observation $y_t$ comes in, from Bayes' rule and the Markov property, we can get $p(x_{t}|y_{1:t})\\propto p(x_t|y_{1:t-1})p(y_{t}|x_t)$, where the first term is the prediction which can be calculated as shown in last exercise and the second term in the new data likelihood for two Gaussian observation models:\n",
    "\n",
    "$$\\text{Posterior}(x_t|y_{1:t}) \\propto \\text{Prediction}(x_t|y_{1:t-1}) \\odot \\text{Likelihood}(y_t|x_t)$$\n",
    "\n",
    "**Suggestions**\n",
    "\n",
    "1. Create a model with switching probability $0.1$ and noise level $0.5$ using function `create_model(switch_prob, noise_level)`\n",
    "\n",
    "2. Complete the code to calculate marginal posterior distribution $p(x_t|y_{1:t-1})$ at time $t$  from last posterior $p(x_{t-1}|y_{1:t-1})$ at time $t-1$ \n",
    "    - Calculate the predictive distrubition $p(x_t =j|y_{1:t-1})=\\sum_i A_{ij} p(x_{t-1}=i|y_{1:t-1})$\n",
    "    - Calculate the likelihood of new data under each component using `exp(model._compute_log_likelihood(yt))`\n",
    "    - Multiply likelihood and prediction element-wise and normalize over two components to get the new posterior probabilities\n",
    "\n",
    "3. Using provided code, plot the average posterior probabilities over time due to evidence together with predictive probabilities without evidence\n",
    "\n",
    "\n",
    "4. Try different values of switching probability (`prob_switch`) and noise level (`noise_level`) either by hand or the widget in section **Interactive Cell** and observe how that affects the posterior probabilities\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "def one_step_update(model, posterior_tm1, Y_t):\n",
    "  \"\"\"Given a HMM model, calculate the one-time-step updates to the posterior.\n",
    "\n",
    "  Args:\n",
    "    model (GaussianHMM instance): the HMM\n",
    "    posterior_tm1 (numpy array): Posterior at `t-1`\n",
    "    Y_t (numpy array): Observation at `t`\n",
    "\n",
    "    Returns:\n",
    "    posterior_t (numpy array): Posterior at `t`\n",
    "  \"\"\"\n",
    "  ##############################################################################\n",
    "  # Insert your code here to:\n",
    "  #      1. Calculate the predicted state given the previous\n",
    "  #      estimate (`posterior_tm1`). Note that `model.transmat_` is equvalent\n",
    "  #      to `A.T`, not `A`.\n",
    "  #      2. Using `model._compute_log_likelihood()`, calculate the likelihood\n",
    "  #      given `Y_t`.\n",
    "  raise NotImplementedError(\"`one_step_update` is incomplete\")\n",
    "  ##############################################################################\n",
    "  prediction = ... @ posterior_tm1\n",
    "  likelihood = np.exp(...)\n",
    "  posterior_t = prediction * likelihood\n",
    "  return posterior_t\n",
    "\n",
    "\n",
    "np.random.seed(101)\n",
    "switch_prob = 0.1\n",
    "noise_level = 0.5\n",
    "nsample = 50\n",
    "T = 160\n",
    "###############################################################################\n",
    "# After finishing the above exercise, please un-comment the following lines\n",
    "###############################################################################\n",
    "# model = create_model(switch_prob, noise_level)\n",
    "\n",
    "# posterior_list = []\n",
    "# for i in range(nsample):\n",
    "#   predictive_probs, posterior_probs = simulate_forward_inference(model, T)\n",
    "#   posterior_list.append(posterior_probs)\n",
    "# posterior_matrix = np.asarray(posterior_list)\n",
    "\n",
    "# plot_evidence_vs_noevidence(posterior_matrix, predictive_probs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 433
    },
    "colab_type": "text",
    "outputId": "66414514-dbdb-4db4-bff2-40d52b34d78a"
   },
   "source": [
    "[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D3_DecisionMaking/solutions/W2D3_Tutorial2_Solution_12fb68dc.py)\n",
    "\n",
    "*Example output:*\n",
    "\n",
    "<img alt='Solution hint' align='left' width=560 height=416 src=https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/tutorials/W2D3_DecisionMaking/static/W2D3_Tutorial2_Solution_12fb68dc_0.png>\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Now you've got a plot of information loss due to diffusion together with the information recovered/uncertainty reduced due to evidence. The difference between the former and the latter is the amount of uncertainty that still remains because of observation noise. You could get a sense of this by changing the noise level in the following cell.\n",
    "\n",
    "**Suggestions**\n",
    "\n",
    "- In the following interactive cell, adjust the switch probability and noise level and observe how information gain changes with signal-to-noise ratio and/or switch probability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 557,
     "referenced_widgets": [
      "52f3c973b57448ea98270c934950296b",
      "e186b20652c34cf7baaa1ebf37e62f01",
      "095f21ab71af4cc4ad5600aa2f81ed4d",
      "653c2902cfe749c8b0aa278036ac0be0",
      "f8a8347274c140dd9401169a7c2dfab4",
      "76f9285d943c49539254ea51a68f7d51",
      "2a19d86baa7d451b96ce3e3a138cc12e",
      "5a63123c482e4aeab29cfcdd193d9d31",
      "f4d17bb5213c4c0da80845b5660cada0",
      "f2e1f1e8541c42bc86753ef1ec512a21",
      "1aa4518cb2614803b12fc2bce85060be",
      "92c257497a264a928de7d6d604618343",
      "337849e02be44791bf7bdfa6080d41c0",
      "e07be1ed09ca4fb0a02a4a1327fe4965",
      "4b184b9f1c7a439ab02cf8ee81cf5919",
      "4c30a9a906a449818b5e3a4a5ca8cf2b"
     ]
    },
    "colab_type": "code",
    "outputId": "aed7b00e-2858-49d1-e93b-88d07f920ac6"
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "\n",
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "np.random.seed(101)\n",
    "@widgets.interact\n",
    "def plot(switch_prob=(0.01, .99, .01), noise_level=(.1, 3, .05),\n",
    "         nsample=(5, 200, 5), T=(20, 300, 5)):\n",
    "  model = create_model(switch_prob, noise_level)\n",
    "  posterior_list = []\n",
    "  for i in range(nsample):\n",
    "    predictive_probs, posterior_probs = simulate_forward_inference(model, T)\n",
    "    posterior_list.append(posterior_probs)\n",
    "  posterior_matrix = np.asarray(posterior_list)\n",
    "  plot_evidence_vs_noevidence(posterior_matrix, predictive_probs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "\n",
    "# Section 3: Inference in a dynamic world"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "\n",
    "\n",
    "### EXERCISE 4: Forward inference of HMM\n",
    "\n",
    "If you set `switch_prob` or `noise_level` to be large in the last exercise, you will observe that some sample inference dots fall below 0.5. This means we are making false inferences about which latent state we are in. \n",
    "\n",
    "In this exercise, let's make a forward inference of a random state sequence rather than a constant one by observing its noisy Gaussian outputs. Different from Exercise 1, here we assume we know the switching probability but don't know the prior (`startprob_`).\n",
    "\n",
    "**Suggestions**\n",
    "\n",
    "1. Build a HMM with prior probabilities= $(0.5,0.5)$, switching probability=$0.1$, and noise level=$1.0$ by calling `create_model(switch_prob, noise_level, startprob)`\n",
    "2. Generate a sample sequence along with observations by calling `model.sample(nstep)`, and use our provided code to visualize the latent trajectory and observations\n",
    "3. Calculate posterior probabilities given data by calling `simulate_forward_inference(model, nstep, observations)`, and make inference of latent states by picking the component with larger posterior probability\n",
    "4. Use our provided code to visualize the inferred state sequence together with the ground truth\n",
    "5. Try different values of switching probability (`prob_switch`) and noise level (`noise_level`) either by hand or the widget in section **Interactive Cell**. When do we start to make false inferences?\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "np.random.seed(101)\n",
    "nstep = 100\n",
    "switch_prob = 0.1\n",
    "log10_noise_level = -1\n",
    "\n",
    "# Build model\n",
    "model = create_model(switch_prob=switch_prob,\n",
    "                     noise_level=10.**log10_noise_level,\n",
    "                     startprob=[0.5, 0.5])\n",
    "observations, states = model.sample(nstep)\n",
    "# Infer state sequence\n",
    "predictive_probs, posterior_probs = simulate_forward_inference(model, nstep,\n",
    "                                                               observations)\n",
    "\n",
    "############################################################################\n",
    "# Insert your code here to:\n",
    "#      Calculate inferred states from posterior probabilities at state 0\n",
    "#      Hint: Compare the probabilities with 0.5 and note that you should\n",
    "#            return 0 if prob > 0.5\n",
    "############################################################################\n",
    "# states_inferred = ...\n",
    "\n",
    "################################################################################\n",
    "# After finishing the above exercises, please un-comment the following lines\n",
    "################################################################################\n",
    "#plot_forward_inference(model, states, observations, states_inferred)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 849
    },
    "colab_type": "text",
    "outputId": "699a4231-9256-4486-b4d7-8ed33c7c7842"
   },
   "source": [
    "[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D3_DecisionMaking/solutions/W2D3_Tutorial2_Solution_94c41b77.py)\n",
    "\n",
    "*Example output:*\n",
    "\n",
    "<img alt='Solution hint' align='left' width=560 height=416 src=https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/tutorials/W2D3_DecisionMaking/static/W2D3_Tutorial2_Solution_94c41b77_0.png>\n",
    "\n",
    "<img alt='Solution hint' align='left' width=560 height=416 src=https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/tutorials/W2D3_DecisionMaking/static/W2D3_Tutorial2_Solution_94c41b77_1.png>\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 905,
     "referenced_widgets": [
      "0e7b3dceac3244159c7527758df3061d",
      "52d6399d321d452a897c92ec4dec8837",
      "d36d85a8faf94e35ac43658149486c2f",
      "87eaf3c4f04f4c02ab5332dd1f13afc9",
      "385446e52edf4ca4bfd69a17051aa939",
      "fd146d2de5db49c7a59b15a135d4a3dc",
      "d9a314dd3a3349c4adf1238c408f95b4",
      "6ede630fecec4bf99013c95afa040690",
      "6eb02d11dd9d4f8a99961bb59bfd59ca",
      "26770d742a5f4a16af2a83f2fe95cb2f"
     ]
    },
    "colab_type": "code",
    "outputId": "99724894-3fde-49bf-9172-d04d1d10cc45"
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "\n",
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "np.random.seed(101)\n",
    "nstep = 100\n",
    "\n",
    "@widgets.interact\n",
    "def plot(switch_prob=(0.01, .99, .01), log10_noise_level=(-8, 1, .01)):\n",
    "  model = create_model(switch_prob=switch_prob,\n",
    "                      noise_level=10.**log10_noise_level,\n",
    "                      startprob = [0.5,0.5])\n",
    "  observations, states = model.sample(nstep)\n",
    "  # observations = observations.flatten()\n",
    "  # Infer state sequence\n",
    "  predictive_probs, posterior_probs = simulate_forward_inference(model, nstep, observations)\n",
    "  states_inferred = posterior_probs <= 0.5\n",
    "  plot_forward_inference(model, states, observations, states_inferred)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "\n",
    "# **Optional** Section: HMM for Poisson spiking neuronal network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 518
    },
    "colab_type": "code",
    "outputId": "9ab4a8cd-5ecc-4367-ef45-5389d6327715"
   },
   "outputs": [],
   "source": [
    "#@title Video 5: HMM for Poisson spiking neurons case study\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"Wb8mf5chmyI\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Given noisy neural or behavioral measurements, we as neuroscientists often want to infer the unobserved latent variables as they change over time. Thalamic relay neurons fire in two distinct modes: a tonic mode where spikes are produced one at a time, and a 'burst mode' where several action potentials are produced in rapid succession. These modes are thought to differentially encode how the neurons relay information from sensory receptors to cortex. A distinct molecular mechanism, T-type calcium channels, switches neurons between modes, but it is very challenging to measure in the brain of a living monkey. However, statistical approaches let us recover the hidden state of those calcium channels purely from their spiking activity, which can be measured in a behaving monkey.\n",
    "\n",
    "Here, we're going to tackle a simplified version of that problem.\n",
    "\n",
    "\n",
    "Let's consider the formulation mentioned in the intro lecture. \n",
    "We have a network of $C$ neurons switching between $K$ states. Neuron $c$ has firing rate $\\lambda_i^c$ in state $i$. The transition between states are represented by the $K\\times K$ transition matrix $A_{ij}$ and initial probability vector $\\psi$ with length $K$ at time $t=1$.\n",
    "\n",
    "Let $y_t^c$ be the number of spikes for cell $c$ in time bin $t$.\n",
    "\n",
    "In the following exercises (4 and 5) and tutorials, you will\n",
    "\n",
    "* Define an instance of such model with $C=5$ and $K=3$\n",
    "* (**Exercise 4**) Generate a dataset from this model\n",
    "* (**Exercise 5**) Implement the M-step for this HMM\n",
    "* Run EM to estimate all parameters $A,\\psi,\\lambda_i^c$\n",
    "* Plot the learning likelihood curve\n",
    "* Plot expected complete log likelihood versus data log likelihood\n",
    "* Compare learnt parameters versus true parameters "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "\n",
    "\n",
    "## **Optional** Section: Define model and generate data\n",
    "\n",
    "Let's first generate a random state sequence from the hidden Markov Chain, and generate `n_frozen_trials` different trials of spike trains for each cell assuming they all use the same underlying sequence we just generated.\n",
    "\n",
    "**Suggestions**\n",
    "\n",
    "1. Run the following two sections **Model and simulation parameters** and **Initialize true model** to define a true model and parameters that will be used in our following exercises. Please take a look at the parameters and come back to these two cells if you encounter a variable you don't know in the future.\n",
    "\n",
    "2. Complete the code to convert a given state sequence to corresponding spike rates for all cells at all times, and use provided code to visualize all spike trains.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {},
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "#@title Helper functions\n",
    "def plot_spike_train(X, Y, dt):\n",
    "    \"\"\"Plots the spike train for cells across trials and overlay the state.\n",
    "\n",
    "      Args:\n",
    "        X: (2d numpy array of binary values): The state sequence in a one-hot\n",
    "                                              representation. (T, states)\n",
    "        Y: (3d numpy array of floats):        The spike sequence.\n",
    "                                              (trials, T, C)\n",
    "        dt (float):                           Interval for a bin.\n",
    "    \"\"\"\n",
    "    n_trials, T, C = Y.shape\n",
    "    trial_T = T * dt\n",
    "    fig = plt.figure(figsize=(.7 * (12.8 + 6.4), .7 * 9.6))\n",
    "\n",
    "    # plot state sequence\n",
    "    starts = [0] + list(np.diff(X.nonzero()[1]).nonzero()[0])\n",
    "    stops = list(np.diff(X.nonzero()[1]).nonzero()[0]) + [T]\n",
    "    states = [X[i + 1].nonzero()[0][0] for i in starts]\n",
    "    for a, b, i in zip(starts, stops, states):\n",
    "        rect = patches.Rectangle((a * dt, 0), (b - a) * dt, n_trials * C,\n",
    "                                 facecolor=plt.get_cmap('tab10').colors[i],\n",
    "                                 alpha=0.15)\n",
    "        plt.gca().add_patch(rect)\n",
    "\n",
    "    # plot rasters\n",
    "    for c in range(C):\n",
    "        if c > 0:\n",
    "            plt.plot([0, trial_T], [c * n_trials, c * n_trials],\n",
    "                     color=plt.get_cmap('tab10').colors[0])\n",
    "        for r in range(n_trials):\n",
    "            tmp = Y[r, :, c].nonzero()[0]\n",
    "            if len(tmp) > 0: plt.plot(np.stack((tmp, tmp)) * dt,\n",
    "                                      (c * n_trials + r + 0.1,\n",
    "                                       c * n_trials + r + .9),\n",
    "                                      'k')\n",
    "\n",
    "    ax = plt.gca()\n",
    "    plt.yticks(np.arange(0, n_trials * C, n_trials),\n",
    "               labels=np.arange(C, dtype=int))\n",
    "    plt.xlabel('time (s)', fontsize=16)\n",
    "    plt.ylabel('Cell number', fontsize=16)\n",
    "\n",
    "def run_em(epochs, Y, psi, A, L, dt):\n",
    "  \"\"\"Run EM for the HMM spiking model.\n",
    "\n",
    "  Args:\n",
    "    epochs (int):       Number of epochs of EM to run\n",
    "    Y (numpy 3d array): Tensor of recordings, has shape (n_trials, T, C)\n",
    "    psi (numpy vector): Initial probabilities for each state\n",
    "    A (numpy matrix):   Transition matrix, A[i,j] represents the prob to switch\n",
    "                        from j to i. Has shape (K,K)\n",
    "    L (numpy matrix):   Poisson rate parameter for different cells.\n",
    "                        Has shape (C,K)\n",
    "    dt (float):         Duration of a time bin\n",
    "\n",
    "  Returns:\n",
    "    save_vals (lists of floats): Data for later plotting\n",
    "    lls (list of flots):         ll Before each EM step\n",
    "    psi (numpy vector):          Estimated initial probabilities for each state\n",
    "    A (numpy matrix):            Estimated transition matrix, A[i,j] represents\n",
    "                                 the prob to switch from j to i. Has shape (K,K)\n",
    "    L (numpy matrix):            Estimated Poisson rate parameter for different\n",
    "                                 cells. Has shape (C,K)\n",
    "  \"\"\"\n",
    "  save_vals = []\n",
    "  lls = []\n",
    "  for e in range(epochs):\n",
    "\n",
    "      # Run E-step\n",
    "      ll, gamma, xi = e_step(Y, psi, A, L, dt)\n",
    "      lls.append(ll)  # log the data log likelihood for current cycle\n",
    "\n",
    "      if e % print_every == 0: print(f'epoch: {e:3d}, ll = {ll}')  # log progress\n",
    "      # Run M-step\n",
    "      psi_new, A_new, L_new = m_step(gamma, xi, dt)\n",
    "\n",
    "      \"\"\"Booking keeping for later plotting\n",
    "      Calculate the difference of parameters for later\n",
    "      interpolation/extrapolation\n",
    "      \"\"\"\n",
    "      dp, dA, dL = psi_new - psi, A_new - A, L_new - L\n",
    "      # Calculate LLs and ECLLs for later plotting\n",
    "      if e in plot_epochs:\n",
    "          b_min = -min([np.min(psi[dp > 0] / dp[dp > 0]),\n",
    "                        np.min(A[dA > 0] / dA[dA > 0]),\n",
    "                        np.min(L[dL > 0] / dL[dL > 0])])\n",
    "          b_max = -max([np.max(psi[dp < 0] / dp[dp < 0]),\n",
    "                        np.max(A[dA < 0] / dA[dA < 0]),\n",
    "                        np.max(L[dL < 0] / dL[dL < 0])])\n",
    "          b_min = np.max([.99 * b_min, b_lims[0]])\n",
    "          b_max = np.min([.99 * b_max, b_lims[1]])\n",
    "          bs = np.linspace(b_min, b_max, num_plot_vals)\n",
    "          bs = sorted(list(set(np.hstack((bs, [0, 1])))))\n",
    "          bs = np.array(bs)\n",
    "          lls_for_plot = []\n",
    "          eclls_for_plot = []\n",
    "          for i, b in enumerate(bs):\n",
    "              ll = e_step(Y, psi + b * dp, A + b * dA, L + b * dL, dt)[0]\n",
    "              lls_for_plot.append(ll)\n",
    "              rate = (L + b * dL) * dt\n",
    "              ecll = ((gamma[:, 0] @ np.log(psi + b * dp) +\n",
    "                       (xi * np.log(A + b * dA)).sum(axis=(-1, -2, -3)) +\n",
    "                       (gamma * stats.poisson(rate).logpmf(Y[..., np.newaxis]).sum(-2)\n",
    "                       ).sum(axis=(-1, -2))).mean() / T / dt)\n",
    "              eclls_for_plot.append(ecll)\n",
    "              if b == 0:\n",
    "                  diff_ll = ll - ecll\n",
    "          lls_for_plot = np.array(lls_for_plot)\n",
    "          eclls_for_plot = np.array(eclls_for_plot) + diff_ll\n",
    "          save_vals.append((bs, lls_for_plot, eclls_for_plot))\n",
    "      # return new parameter\n",
    "      psi, A, L = psi_new, A_new, L_new\n",
    "\n",
    "  ll = e_step(Y, psi, A, L, dt)[0]\n",
    "  lls.append(ll)\n",
    "  print(f'epoch: {epochs:3d}, ll = {ll}')\n",
    "  return save_vals, lls, psi, A, L\n",
    "\n",
    "def plot_lls(lls):\n",
    "  \"\"\"Plots log likelihoods at each epoch.\n",
    "  Args:\n",
    "    lls (list of floats) log likelihoods at each epoch.\n",
    "  \"\"\"\n",
    "  epochs = len(lls)\n",
    "  fig, ax = plt.subplots()\n",
    "  ax.plot(range(epochs) , lls, linewidth=3)\n",
    "  span = max(lls) - min(lls)\n",
    "  ax.set_ylim(min(lls) - span * 0.05, max(lls) + span * 0.05)\n",
    "  plt.xlabel('iteration')\n",
    "  plt.ylabel('log likelihood')\n",
    "  plt.show(fig)\n",
    "\n",
    "def plot_lls_eclls(plot_epochs, save_vals):\n",
    "  \"\"\"Plots log likelihoods at each epoch.\n",
    "  Args:\n",
    "    plot_epochs (list of ints):  Which epochs were saved to plot.\n",
    "    save_vals (lists of floats): Different likelihoods from EM for plotting.\n",
    "  \"\"\"\n",
    "  rows = int(np.ceil(min(len(plot_epochs), len(save_vals)) / 3))\n",
    "  fig, axes = plt.subplots(rows, 3, figsize=(.7 * 6.4 * 3, .7 * 4.8 * rows))\n",
    "  axes = axes.flatten()\n",
    "\n",
    "  minll, maxll = np.inf, -np.inf\n",
    "  for i, (ax, (bs, lls_for_plot, eclls_for_plot)) in enumerate(zip(axes, save_vals)):\n",
    "      ax.set_xlim([-1.15, 2.15])\n",
    "      min_val = np.stack((lls_for_plot, eclls_for_plot)).min()\n",
    "      max_val = np.stack((lls_for_plot, eclls_for_plot)).max()\n",
    "\n",
    "      ax.plot([0, 0], [min_val, lls_for_plot[bs == 0]], '--b')\n",
    "      ax.plot([1, 1], [min_val, lls_for_plot[bs == 1]], '--b')\n",
    "      ax.set_xticks([0, 1])\n",
    "      ax.set_xticklabels([f'$\\\\theta^{plot_epochs[i]}$',\n",
    "                          f'$\\\\theta^{plot_epochs[i] + 1}$'])\n",
    "      ax.tick_params(axis='y')\n",
    "      ax.tick_params(axis='x')\n",
    "\n",
    "\n",
    "      ax.plot(bs, lls_for_plot)\n",
    "      ax.plot(bs, eclls_for_plot)\n",
    "\n",
    "      if min_val < minll: minll = min_val\n",
    "      if max_val > maxll: maxll = max_val\n",
    "\n",
    "      if i % 3 == 0: ax.set_ylabel('log likelihood')\n",
    "      if i == 4:\n",
    "          l = ax.legend(ax.lines[-2:], ['LL', 'ECLL'], framealpha=1)\n",
    "  plt.show(fig)\n",
    "\n",
    "def plot_learnt_vs_true(L_true, L, A_true, A, dt):\n",
    "  \"\"\"Plot and compare the true and learnt parameters.\n",
    "\n",
    "  Args:\n",
    "    L_true (numpy array): True L.\n",
    "    L (numpy array):      Estimated L.\n",
    "    A_true (numpy array): True A.\n",
    "    A (numpy array):      Estimated A.\n",
    "    dt (float):           Bin length.\n",
    "  \"\"\"\n",
    "  C, K = L.shape\n",
    "  fig = plt.figure(figsize=(8, 4))\n",
    "  plt.subplot(121)\n",
    "  plt.plot([0, L_true.max() * 1.05], [0, L_true.max() * 1.05], '--b')\n",
    "  for i in range(K):\n",
    "      for c in range(C):\n",
    "          plt.plot(L_true[c, i], L[c, i], color='C{}'.format(c),\n",
    "                   marker=['o', '*', 'd'][i])  # this line will fail for K > 3\n",
    "  ax = plt.gca()\n",
    "  ax.axis('equal')\n",
    "  plt.xlabel('True firing rate (Hz)')\n",
    "  plt.ylabel('Inferred firing rate (Hz)')\n",
    "  xlim, ylim = ax.get_xlim(), ax.get_ylim()\n",
    "  for c in range(C):\n",
    "      plt.plot([-10^6], [-10^6], 'o', color='C{}'.format(c))\n",
    "  for i in range(K):\n",
    "      plt.plot([-10^6], [-10^6], '.', marker=['o', '*', 'd'][i], c=\"black\")\n",
    "  l = plt.legend(ax.lines[-C - K:],\n",
    "                 [f'cell {c + 1}' for c in range(C)] + [f'state {i + 1}' for i in range(K)])\n",
    "  ax.set_xlim(xlim), ax.set_ylim(ylim)\n",
    "\n",
    "  plt.subplot(122)\n",
    "  ymax = np.max(A_true - np.diag(np.diag(A_true))) / dt * 1.05\n",
    "  plt.plot([0, ymax], [0, ymax], '--b')\n",
    "  for j in range(K):\n",
    "      for i in range(K):\n",
    "          if i == j: continue\n",
    "          plt.plot(A_true[i, j] / dt, A[i, j] / dt, 'o')\n",
    "  ax = plt.gca()\n",
    "  ax.axis('equal')\n",
    "  plt.xlabel('True transition rate (Hz)')\n",
    "  plt.ylabel('Inferred transition rate (Hz)')\n",
    "  l = plt.legend(ax.lines[1:], ['state 1 -> 2',\n",
    "                                      'state 1 -> 3',\n",
    "                                      'state 2 -> 1',\n",
    "                                      'state 2 -> 3',\n",
    "                                      'state 3 -> 1',\n",
    "                                      'state 3 -> 2'\n",
    "                                    ])\n",
    "  plt.show(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### Model and simulation parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "# model and data parameters\n",
    "C = 5  # number of cells\n",
    "K = 3  # number of states\n",
    "dt = 0.002  # seconds\n",
    "trial_T = 2.0  # seconds\n",
    "n_frozen_trials = 20  # used to plot multiple trials with the same state sequence\n",
    "n_trials = 300  # number of trials (each has it's own state sequence)\n",
    "\n",
    "# for random data\n",
    "max_firing_rate = 50  # Hz\n",
    "max_transition_rate = 3  # Hz\n",
    "\n",
    "# needed to plot LL and ECLL for every M-step\n",
    "# **This substantially slows things down!!**\n",
    "num_plot_vals = 10  # resolution of the plot (this is the expensive part)\n",
    "b_lims = (-1, 2)  # lower limit on graph (b = 0 is start-of-M-step LL; b = 1 is end-of-M-step LL)\n",
    "plot_epochs = list(range(9))  # list of epochs to plot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### Initialize true model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "np.random.seed(101)\n",
    "T = round(trial_T / dt)\n",
    "ts = np.arange(T)\n",
    "\n",
    "# initial state distribution\n",
    "psi = np.arange(1, K + 1)\n",
    "psi = psi / psi.sum()\n",
    "\n",
    "# off-diagonal transition rates sampled uniformly\n",
    "A = np.random.rand(K, K) * max_transition_rate * dt\n",
    "A = (1. - np.eye(K)) * A\n",
    "A = A + np.diag(1 - A.sum(1))\n",
    "\n",
    "# hand-crafted firing rates make good plots\n",
    "L = np.array([\n",
    "    [.02, .8, .37],\n",
    "    [1., .7, .1],\n",
    "    [.92, .07, .5],\n",
    "    [.25, .42, .75],\n",
    "    [.15, .2, .85]\n",
    "]) * max_firing_rate     # (C,K)\n",
    "\n",
    "# Save true parameters for comparison later\n",
    "psi_true = psi\n",
    "A_true = A\n",
    "L_true = L"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### Generate data with frozen sequence and plot\n",
    "Given a state sequence `[0,1,1,3,2,...]`, we'll first convert each state in to sequence in to the so-called \"one-hot\" coding. For example, with 5 total states, the one-hot coding of state `0` is `[1,0,0,0,0]` and the coding for state `3` is `[0,0,0,1,0]`. Suppose we now have a sequence of length `T`, the one-hot coding of this sequence `Xf` will have shape `(T,K)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 481
    },
    "colab_type": "code",
    "outputId": "917adc3e-7634-4ea8-8b62-ab8c9dead2f2"
   },
   "outputs": [],
   "source": [
    "np.random.seed(101)\n",
    "# sample n_frozen_trials state sequences\n",
    "Xf = np.zeros(T, dtype=int)\n",
    "Xf[0] = (psi.cumsum() > np.random.rand()).argmax()\n",
    "for t in range(1, T):\n",
    "    Xf[t] = (A[Xf[t - 1],:].cumsum() > np.random.rand()).argmax()\n",
    "\n",
    "# switch to one-hot encoding of the state\n",
    "Xf = np.eye(K, dtype=int)[Xf]  # (T,K)\n",
    "\n",
    "# get the Y values\n",
    "Rates = np.squeeze(L @ Xf[..., None]) * dt  # (T,C)\n",
    "\n",
    "Rates = np.tile(Rates, [n_frozen_trials, 1, 1]) # (n_trials, T, C)\n",
    "Yf = stats.poisson(Rates).rvs()\n",
    "\n",
    "with plt.xkcd():\n",
    "  plot_spike_train(Xf, Yf, dt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### Generate data for EM learning\n",
    "\n",
    "The previous dataset is generated with the same state sequence for visualization. Now let's generate `n_trials` trials of observations, each one with its own randomly generated sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "outputId": "56bd4317-8cac-4e7e-9ddc-bc7816f63d99"
   },
   "outputs": [],
   "source": [
    "np.random.seed(101)\n",
    "# sample n_trials state sequences\n",
    "X = np.zeros((n_trials, T), dtype=int)\n",
    "X[:, 0] = (psi_true.cumsum(0)[:, None] > np.random.rand(n_trials)).argmax(0)\n",
    "for t in range(1, T):\n",
    "  X[:, t] = (A_true[X[:, t - 1], :].T.cumsum(0) > np.random.rand(n_trials)).argmax(0)\n",
    "\n",
    "# switch to one-hot encoding of the state\n",
    "one_hot = np.eye(K)[np.array(X).reshape(-1)]\n",
    "X = one_hot.reshape(list(X.shape) + [K])\n",
    "\n",
    "# get the Y values\n",
    "Y = stats.poisson(np.squeeze(L_true @ X[..., None]) * dt).rvs()  # (n_trials, T, C)\n",
    "print(\"Y has shape: (n_trial={},T={},C={})\".format(*Y.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "# ADVICE FROM NMA  \n",
    "We, the organizers, know that the next section is much longer and more challenging than most other tutorial content. **We do not expect you to finish it—or even start it—right now**. In fact, we strongly suggest saving your time and energy for the Kalman Filtering introduced in Tutorial 3, because it will play an important role in tomorrow's material too. Since Kalman filters also use the E-M algorithm, learning about them first may provide you with useful intuitions about this material. \n",
    "\n",
    "That said, the EM algorithm can be a very useful and powerful optimization tool. Since it is typically taught in the context of Hidden Markov Models, we have included it here for your reference. \n",
    "\n",
    "To reiterate, the remainder of this notebook is *completely* and *absolutely* optional. It is not essential to understand the rest of the NMA content. By this point in Tutorial 2, we believe that you will have seen enough about HMMs to know when/if they might be relevant for your own research. When that day comes, or you are just feeling curious, this material will be here waiting! "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "\n",
    "## **Optional** Section: EM algorithm for HMM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 518
    },
    "colab_type": "code",
    "outputId": "2cf578fb-805f-4758-bc4b-de604eaeea6d"
   },
   "outputs": [],
   "source": [
    "#@title Video 6: EM Tutorial\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"umU4wUWlKvg\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Finding the optimal values of parameters that maximizes the data likelihood is practically infeasible since we need to integrating out all latent variables $x_{1:T}$. The time needed is exponential to $T$. Thus as an alternative approach, we use the Expectation-Maximization algorithm, which iteratively performing a E-step followed by a M-step and is guaranteed to not decrease(usually increase) the data likelihood after each EM cycle.\n",
    "\n",
    "\n",
    "In this section we will briefly review the EM algorithm  for HMM and list\n",
    "\n",
    "* Recursive equations for forward and backward probabilities $a_i(t)$ and $b_i(t)$\n",
    "* Expressions for singleton and pairwise marginal distributions after seeing data: $\\gamma_{i}(t):=p_{\\theta}\\left(x_{t}=i | Y_{1: T}\\right)$ and $\\xi_{i j}(t) = p_{\\theta}(x_t=i,x_{t+1}=j|Y_{1:T})$\n",
    "* Closed-form solutions for updated values of $A,\\psi,\\lambda$ which increases data likelihood\n",
    "\n",
    "\n",
    "---\n",
    "#### E-step: Forward-backward algorithm \n",
    "In the forward pass, we calculate the **forward probabilities**, or the joint probability of $x_t$ and current and past data $Y_{1:t}$: $a_i(t):=p(x_t=i,Y_{1:t})$ recursively by\n",
    "\n",
    "$$a_i(t) = p_(y_t|x_i=t)\\sum_j A_{ji} a_j(t-1)$$\n",
    "\n",
    "In contrast to the intro, now $A_{ji}$ means **the transition probability from state $j$ to state $i$.**\n",
    "\n",
    "The backward pass calculate the **backward probabilities** $b_i(t):=p_{\\theta}(Y_{t+1:T}|x_t=i)$, which is the likelihood of observing all future data points given current state $x_t$. The recursion of $b_i(t)$ is given by \n",
    "\n",
    "$$ b_i(t) = \\sum_j p_{\\theta}(y_{t+1}|x_{t+1}=j)b_j(t+1)A_{ij} $$\n",
    "\n",
    "Combining all past and future information, the **singleton and pairwise marginal distributions** are given by \n",
    "\n",
    "$$ \\gamma_{i}(t):=p_{\\theta}\\left(x_{t}=i | Y_{1: T}\\right)=\\frac{a_{i}(t) b_{i}(t)}{p_{\\theta}\\left(Y_{1: T}\\right)} $$ \n",
    "\n",
    "$$  \\xi_{i j}(t) = p_{\\theta}(x_t=i,x_{t+1}=j|Y_{1:T}) =\\frac{b_{j}(t+1)p_{\\theta}\\left(y_{t+1} | x_{t+1}=j\\right) A_{i j} a_{i}(t)}{p_{\\theta}\\left(Y_{1: T}\\right)} $$\n",
    "\n",
    "where $p_{\\theta}(Y_{1:T})=\\sum_i a_i(T)$.\n",
    "\n",
    "---\n",
    "#### M-step\n",
    "\n",
    "The M-step for HMM has a closed-form solution. First the new transition matrix is given by \n",
    "$$ \n",
    " A_{ij} =\\frac{\\sum_{t=1}^{T-1} \\xi_{i j}(t)}{\\sum_{t=1}^{T-1} \\gamma_{i}(t)}\n",
    "$$\n",
    "\n",
    "which is the expected empirical transition probabilities. \n",
    "New initial probabilities and parameters of the emission models are also given by their empirical values given single and pairwise marginal distributions:\n",
    "\n",
    "$$ \\psi_i = \\frac{1}{N}\\sum_{trials}\\gamma_i(1) $$ \n",
    "\n",
    "$$ \\lambda_{i}^{c}=\\frac{\\sum_{t} \\gamma_{i}(t) y_{t}^{c}}{\\sum_{t} \\gamma_{i}(t) d t}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "\n",
    "\n",
    "### E-step: forward and backward algorithm\n",
    "\n",
    "**(Optional)**\n",
    "\n",
    "In this section you will read through the code for the forward-backward algorithm and understand how to implement the computation efficiently in `numpy` by calculating the recursion for all trials at once.\n",
    "\n",
    "---\n",
    "\n",
    "Let's re-write the forward and backward recursions in a more compact form:\n",
    "\n",
    "$$ a_i^t = \\sum_j A_{ji}o_j^t a_j^{t-1}  $$\n",
    "\n",
    "\n",
    "$$b^t_i = \\sum_j A_{ij} o_j^{t+1}b_j^{t+1} $$ where $o_j^{t}=p(y_{t}|x_{t}=j)$.\n",
    "\n",
    "\n",
    "Let's take the backward recursion for example. In practice we will handle all trials together since they are independent of each other. After adding a trial index $l$ to the recursion equations, the backward recursion becomes:\n",
    "\n",
    "$$b^t_{li} = \\sum_j A_{ij} o_{lj}^{t+1}b_{lj}^{t+1} $$ \n",
    "\n",
    "What we have in hand are:\n",
    "* `A`: matrix of size `(K,K)`\n",
    "* `o^{t+1}`: array of size `(N,K)` is the log data likelihood for all trials at a given time\n",
    "* `b^{t+1}`: array of size `(N,K)` is the backward probability for all trials at a given time\n",
    "\n",
    "where `N` stands for the number of trials.\n",
    "\n",
    "The index size and meaning doesn't match for these three arrays: the index is $i$ for $A$ in the first dimension and is $l$ for $o$ and $b$, so we can't just multiply them together. However, we can do this by viewing vectors $o^{t+1}_{l\\cdot}$ and $b^{t+1}_{l\\cdot}$ as a matrix with 1 row and re-write the backward equation as:\n",
    "\n",
    "$$b^t_{li} = \\sum_j A_{ij} o_{l1j}^{t+1}b_{l1j}^{t+1} $$\n",
    "\n",
    "Now we can just multiply these three arrays element-wise and sum over the last dimension.\n",
    "\n",
    "In `numpy`, we can achieve this by indexing the array with `None` at the location we want to insert a dimension. Take `b` with size `(N,T,K)` for example,`b[:,t,:]` will have shape `(N,K)`, `b[:,t,None,:]` will have shape `(N,1,K)` and `b[:,t,:,None]` will have shape `(N,K,1)`.\n",
    "\n",
    "So the backward recursion computation can be implemented as \n",
    "\n",
    "```python\n",
    "b[:,t,:] = (A * o[:,t+1,None,:] * b[:,t+1,None,:]).sum(-1)\n",
    "```\n",
    "\n",
    "---\n",
    "\n",
    "In addition to the trick introduced above, in this exercise we will work in the **log scale** for numerical stability.\n",
    "\n",
    "\n",
    "**Suggestions** \n",
    "\n",
    "1. Take a look at the code for the forward recursion and backward recursion.\n",
    "\n",
    "    \n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "def e_step(Y, psi, A, L, dt):\n",
    "  \"\"\"Calculate the E-step for the HMM spiking model.\n",
    "\n",
    "  Args:\n",
    "    Y (numpy 3d array): tensor of recordings, has shape (n_trials, T, C)\n",
    "    psi (numpy vector): initial probabilities for each state\n",
    "    A (numpy matrix):   transition matrix, A[i,j] represents the prob to\n",
    "                        switch from i to j. Has shape (K,K)\n",
    "    L (numpy matrix):   Poisson rate parameter for different cells.\n",
    "                        Has shape (C,K)\n",
    "    dt (float):         Bin length\n",
    "\n",
    "  Returns:\n",
    "    ll (float):             data log likelihood\n",
    "    gamma (numpy 3d array): singleton marginal distribution.\n",
    "                            Has shape (n_trials, T, K)\n",
    "    xi (numpy 4d array):    pairwise marginal distribution for adjacent\n",
    "                            nodes . Has shape (n_trials, T-1, K, K)\n",
    "  \"\"\"\n",
    "  n_trials = Y.shape[0]\n",
    "  T = Y.shape[1]\n",
    "  K = psi.size\n",
    "  log_a = np.zeros((n_trials, T, K))\n",
    "  log_b = np.zeros((n_trials, T, K))\n",
    "\n",
    "  log_A = np.log(A)\n",
    "  log_obs = stats.poisson(L * dt).logpmf(Y[..., None]).sum(-2)  # n_trials, T, K\n",
    "\n",
    "  # forward pass\n",
    "  log_a[:, 0] = log_obs[:, 0] + np.log(psi)\n",
    "  for t in range(1, T):\n",
    "    tmp = log_A + log_a[:, t - 1, : ,None]  # (n_trials, K,K)\n",
    "    maxtmp = tmp.max(-2)  # (n_trials,K)\n",
    "    log_a[:, t] = (log_obs[:, t] + maxtmp +\n",
    "                    np.log(np.exp(tmp - maxtmp[:, None]).sum(-2)))\n",
    "\n",
    "  # backward pass\n",
    "  for t in range(T - 2, -1, -1):\n",
    "    tmp = log_A + log_b[:, t + 1, None] + log_obs[:, t + 1, None]\n",
    "    maxtmp = tmp.max(-1)\n",
    "    log_b[:, t] = maxtmp + np.log(np.exp(tmp - maxtmp[..., None]).sum(-1))\n",
    "\n",
    "  # data log likelihood\n",
    "  maxtmp = log_a[:, -1].max(-1)\n",
    "  ll = np.log(np.exp(log_a[:, -1] - maxtmp[:, None]).sum(-1)) + maxtmp\n",
    "\n",
    "  # singleton and pairwise marginal distributions\n",
    "  gamma = np.exp(log_a + log_b - ll[:, None, None])\n",
    "  xi = np.exp(log_a[:, :-1, :, None] + (log_obs + log_b)[:, 1:, None] +\n",
    "              log_A - ll[:, None, None, None])\n",
    "\n",
    "  return ll.mean() / T / dt, gamma, xi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 518
    },
    "colab_type": "code",
    "outputId": "47819702-3e27-4afd-e5e5-2ea98662ebf0"
   },
   "outputs": [],
   "source": [
    "#@title EXERCISE 7: Implement the M-step Video\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"H4GGTg_9BaE\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "\n",
    "### EXERCISE 5: Implement the M-step\n",
    "\n",
    "In this exercise you will complete the M-step for this HMM using closed form solutions mentioned before.\n",
    "\n",
    "**Suggestions**\n",
    "\n",
    "1. Calculate new initial probabilities as empirical counts of singleton marginals\n",
    "\n",
    "$$ \\psi_i = \\frac{1}{N}\\sum_{trials}\\gamma_i(1) $$\n",
    "\n",
    "2. Remember the extra trial dimension and average over all trials\n",
    "\n",
    "\n",
    "**For reference:**\n",
    "\n",
    "New transition matrix is calculated as empirical counts of transition events from marginals\n",
    "\n",
    "$$ A_{ij} =\\frac{\\sum_{t=1}^{T-1} \\xi_{i j}(t)}{\\sum_{t=1}^{T-1} \\gamma_{i}(t)}$$ \n",
    "\n",
    "\n",
    "New spiking rates for each cell and each state are given by\n",
    "\n",
    "$$  \\lambda_{i}^{c}=\\frac{\\sum_{t} \\gamma_{i}(t) y_{t}^{c}}{\\sum_{t} \\gamma_{i}(t) d t} $$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "def m_step(gamma, xi, dt):\n",
    "  \"\"\"Calculate the M-step updates for the HMM spiking model.\n",
    "\n",
    "  Args:\n",
    "    gamma ():       Number of epochs of EM to run\n",
    "    xi (numpy 3d array): Tensor of recordings, has shape (n_trials, T, C)\n",
    "    dt (float):         Duration of a time bin\n",
    "\n",
    "  Returns:\n",
    "    psi_new (numpy vector): Updated initial probabilities for each state\n",
    "    A_new (numpy matrix):   Updated transition matrix, A[i,j] represents the\n",
    "                            prob. to switch from j to i. Has shape (K,K)\n",
    "    L_new (numpy matrix):   Updated Poisson rate parameter for different\n",
    "                            cells. Has shape (C,K)\n",
    "  \"\"\"\n",
    "  raise NotImplementedError(\"`m_step` need to be implemented\")\n",
    "  ############################################################################\n",
    "  # Insert your code here to:\n",
    "  #    Calculate the new prior probabilities in each state at time 0\n",
    "  #    Hint: Take the first time step and average over all trials\n",
    "  ###########################################################################\n",
    "  psi_new = ...\n",
    "  # Make sure the probabilities are normalized\n",
    "  psi_new /= psi_new.sum()\n",
    "\n",
    "  # Calculate new transition matrix\n",
    "  A_new = xi.sum(axis=(0, 1)) / gamma[:, :-1].sum(axis=(0, 1))[:, np.newaxis]\n",
    "  # Calculate new firing rates\n",
    "  L_new = (np.swapaxes(Y, -1, -2) @ gamma).sum(axis=0) / gamma.sum(axis=(0, 1)) / dt\n",
    "  return psi_new, A_new, L_new"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {},
    "colab_type": "text"
   },
   "source": [
    "[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D3_DecisionMaking/solutions/W2D3_Tutorial2_Solution_099bf548.py)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 518
    },
    "colab_type": "code",
    "outputId": "86db5e5a-e853-4b88-e7f1-737d32a656f5"
   },
   "outputs": [],
   "source": [
    "#@title Video 8: Running and plotting EM\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"6UTsXxE3hG0\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "\n",
    "### Run EM\n",
    "\n",
    "####Initialization for parameters\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "np.random.seed(101)\n",
    "# number of EM steps\n",
    "epochs = 9\n",
    "print_every = 1\n",
    "\n",
    "# initial state distribution\n",
    "psi = np.arange(1, K + 1)\n",
    "psi = psi / psi.sum()\n",
    "\n",
    "# off-diagonal transition rates sampled uniformly\n",
    "A = np.ones((K, K)) * max_transition_rate * dt / 2\n",
    "A = (1 - np.eye(K)) * A\n",
    "A = A + np.diag(1 - A.sum(1))\n",
    "\n",
    "# firing rates sampled uniformly\n",
    "L = np.random.rand(C, K) * max_firing_rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 242
    },
    "colab_type": "code",
    "outputId": "29b2b875-a012-4eb2-e50a-8b140acf2aa2"
   },
   "outputs": [],
   "source": [
    "# LL for true vs. initial parameters\n",
    "print(f'LL for true 𝜃:    {e_step(Y, psi_true, A_true, L_true, dt)[0]}')\n",
    "print(f'LL for initial 𝜃: {e_step(Y, psi, A, L, dt)[0]}\\n')\n",
    "\n",
    "# Run EM\n",
    "save_vals, lls, psi, A, L = run_em(epochs, Y, psi, A, L, dt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "# EM doesn't guarantee the order of learnt latent states are the same as that of true model\n",
    "# so we need to sort learnt parameters\n",
    "\n",
    "# Compare all true and estimated latents across cells\n",
    "cost_mat = np.sum((L_true[..., np.newaxis] - L[:, np.newaxis])**2, axis=0)\n",
    "true_ind, est_ind = linear_sum_assignment(cost_mat)\n",
    "\n",
    "psi = psi[est_ind]\n",
    "A = A[est_ind]\n",
    "A = A[:, est_ind]\n",
    "L = L[:, est_ind]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "## **Optional** Section: Plotting the training process and learnt model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### Plotting progress during EM!\n",
    "\n",
    "Now you can\n",
    "\n",
    "* Plot the likelihood during training\n",
    "* Plot the M-step log likelihood versus expected complete log likelihood(ECLL) to get an intuition of how EM works and the convexity of ECLL\n",
    "* Plot learnt parameters versus true parameters "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 433
    },
    "colab_type": "code",
    "outputId": "f51c6cad-2d5d-4113-9825-370072bdc90f"
   },
   "outputs": [],
   "source": [
    "# Plot the log likelihood after each epoch of EM\n",
    "with plt.xkcd():\n",
    "  plot_lls(lls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 725
    },
    "colab_type": "code",
    "outputId": "2ec5ba79-236f-4627-d596-6f8f52016ce2"
   },
   "outputs": [],
   "source": [
    "# For each saved epoch, plot the log likelihood and expected complete log likelihood\n",
    "# for the initial and final parameter values\n",
    "with plt.xkcd():\n",
    "  plot_lls_eclls(plot_epochs, save_vals)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### Plot learnt parameters vs. true parameters\n",
    "\n",
    "Now we will plot the (sorted) learnt parameters with true parameters to see if we successfully recovered all the parameters\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 288
    },
    "colab_type": "code",
    "outputId": "b2269d42-30f9-492f-8831-74cba676f134"
   },
   "outputs": [],
   "source": [
    "# Compare true and learnt parameters\n",
    "with plt.xkcd():\n",
    "  plot_learnt_vs_true(L_true, L, A_true, A, dt)"
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "W2D3_Tutorial2",
   "provenance": [],
   "toc_visible": true
  },
  "kernel": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.8"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": true,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
